/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.webank.tikv;

import com.google.protobuf.ByteString;
import com.webank.tikv.common.TiTimestamp;
import com.webank.tikv.exception.TiKVException;
import com.webank.tikv.kvproto.Errorpb;
import com.webank.tikv.kvproto.Metapb;
import com.webank.tikv.kvproto.Pdpb;
import com.webank.tikv.region.RegionManager;
import com.webank.tikv.region.RegionStoreClient;
import com.webank.tikv.region.TiRegion;
import com.webank.tikv.region.TiStore;
import com.webank.tikv.util.BackOffer;
import com.webank.tikv.util.ChannelFactory;
import com.webank.tikv.util.ConcreteBackOffer;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * TiSession is the holder for PD Client, Store pdClient and PD Cache All sessions share common
 * region store connection pool but separated PD conn and cache for better concurrency
 *
 * <p>TiSession is thread-safe but it's also recommended to have multiple session avoiding lock
 * contention
 */
public class TiSession {

    private static final Logger LOGGER = LogManager.getLogger(TiSession.class);

    private final TiConfiguration conf;
    private final ChannelFactory channelFactory;
    private final RegionStoreClient.RegionStoreClientBuilder clientBuilder;

    private final PDClient pdClient;
    private final RegionManager regionManager;

    private final AtomicBoolean isClosed = new AtomicBoolean();

    public static TiSession create(TiConfiguration conf) {
        return new TiSession(conf);
    }

    private TiSession(TiConfiguration conf) {
        this.conf = conf;
        if (conf.isTlsEnable()) {
            channelFactory = new ChannelFactory(conf.getMaxFrameSize(), conf.getKeepAliveTime(),
                conf.getKeepAliveTimeout(), conf.getIdleTimeout(), conf.getConnRecycleTime(),
                conf.getCertReloadInterval(), conf.getTrustCertCollectionFile(),
                conf.getKeyCertChainFile(), conf.getKeyFile());
        } else if (conf.isJksEnable()) {
            channelFactory = new ChannelFactory(conf.getMaxFrameSize(), conf.getKeepAliveTime(),
                conf.getKeepAliveTimeout(), conf.getIdleTimeout(), conf.getConnRecycleTime(),
                conf.getCertReloadInterval(), conf.getJksKeyPath(), conf.getJksKeyPassword(),
                conf.getJksTrustPath(), conf.getJksTrustPassword());
        } else {
            channelFactory = new ChannelFactory(conf.getMaxFrameSize(), conf.getKeepAliveTime(),
                conf.getKeepAliveTimeout(), conf.getIdleTimeout(), conf.getConnRecycleTime());
        }

        this.pdClient = new PDClient(conf, channelFactory);
        this.regionManager = new RegionManager(conf, pdClient, channelFactory);
        this.clientBuilder = new RegionStoreClient.RegionStoreClientBuilder(conf, channelFactory, regionManager);

        if (conf.isWarmUpEnable()) {
            warmUp();
        }
        LOGGER.info("TiSession initialized");
    }

    public void warmUp() {
        long warmUpStartTime = System.nanoTime();
        BackOffer backOffer = ConcreteBackOffer.newCustomBackOff(5000);
        try {
            // let JVM ClassLoader load gRPC error related classes
            // this operation may cost 100ms
            Errorpb.Error.newBuilder().setNotLeader(Errorpb.NotLeader.newBuilder().build()).build();

            // warm up store cache
            List<Metapb.Store> stores = pdClient.getAllStores(backOffer);
            for (Metapb.Store store : stores) {
                regionManager.updateStore(null, new TiStore(store));
            }

            // use scan region to load region cache with limit
            ByteString startKey = ByteString.EMPTY;
            do {
                List<Pdpb.Region> regions = pdClient.scanRegions(backOffer, startKey,
                    ByteString.EMPTY, conf.getScanRegionsLimit());
                if (regions == null || regions.isEmpty()) {
                    // something went wrong, but the warm-up process could continue
                    break;
                }
                for (Pdpb.Region region : regions) {
                    TiRegion tiRegion = regionManager.createRegion(region.getRegion(), region.getLeader(), backOffer);
                    regionManager.putRegion(tiRegion);
                }
                startKey = regions.get(regions.size() - 1).getRegion().getEndKey();
            } while (!startKey.isEmpty());

        } catch (Exception e) {
            LOGGER.warn("warm up fails:{}", e.getMessage());
        } finally {
            LOGGER.info("warm up success duration:{}ms", (System.nanoTime() - warmUpStartTime) / 1_000_000);
        }
    }

    public RawKVClient createRawClient() {
        checkIsClosed();
        return new RawKVClient(this, clientBuilder);
    }

    public TiTimestamp getTimestamp() {
        checkIsClosed();
        return pdClient.getTimestamp(ConcreteBackOffer.newCustomBackOff(BackOffer.PD_INFO_BACKOFF));
    }

    public RegionManager getRegionManager() {
        checkIsClosed();
        return regionManager;
    }

    public PDClient getPdClient() {
        checkIsClosed();
        return pdClient;
    }

    public TiConfiguration getConf() {
        return conf;
    }

    public boolean isClosed() {
        return isClosed.get();
    }

    private void checkIsClosed() {
        if (isClosed.get()) {
            throw new TiKVException("this TiSession is closed");
        }
    }

    public void close() {
        if (!isClosed.compareAndSet(false, true)) {
            return;
        }

        if (regionManager != null) {
            regionManager.close();
        }
        if (pdClient != null) {
            pdClient.close();
        }
        if (channelFactory != null) {
            channelFactory.close();
        }
    }

}
